{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Code for deep recurrent Q learning\n",
    "# Paper: Deep Recurrent Q-Learning for Partially Observable MDPs\n",
    "# Paper url: https://arxiv.org/pdf/1507.06527.pdf\n",
    "\n",
    "import gym, random, pickle, os.path, math, glob\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.autograd as autograd\n",
    "import pdb\n",
    "import os\n",
    "\n",
    "from atari_wrappers import make_atari, wrap_deepmind,LazyFrames\n",
    "from IPython.display import clear_output\n",
    "from tensorboardX import SummaryWriter\n",
    "\n",
    "USE_CUDA = torch.cuda.is_available()\n",
    "dtype = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor\n",
    "Variable = lambda *args, **kwargs: autograd.Variable(*args, **kwargs).cuda() if USE_CUDA else autograd.Variable(*args, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x1f39dc79688>"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD7CAYAAACscuKmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8vihELAAAACXBIWXMAAAsTAAALEwEAmpwYAAASR0lEQVR4nO3da4wdZ33H8e9v5pyza6/tOM6FJnaUBAkFECgJdSE0VQUEo0AjqCqBSEWFKtS8oW2okChpXyDe+UWFQCqlsrg0Ldc0QEERAiIuopVQmpCkFHBCuKSJc7MxhNje7O45M/++mHG8mF3v+Ozu7Jl9fh9ptTtzzp55np39nZkzl+eviMDMNr9soxtgZu1w2M0S4bCbJcJhN0uEw26WCIfdLBGrCruk6yU9KOknkt67Vo0ys7Wncc+zS8qBHwP7gEPA3cCNEfGjtWuema2V3ip+9+XATyLiZwCSPgu8CVg27ANNxTQzK76wBn2il4MEWkUL14pEAGRQ9EXkY75MCdkQFAFl/d3Wx6J1VvZFmcg6m5t7moXhiSVTs5qw7wYeXTR9CHjFmX5hmhleoevO+KLq9ch3X0xx3nYizygHOeQbm/jIRJmLcpBx/OIe8+eO157+8WDmqZJ8riRbKMmH5Rq31E76jXV20Rqss/lurLO77/nwso+tJuxL/fV+621P0k3ATQDT+XZ6u/ec+VXzjHLHVspBr1qCNn7THhJkIjIx2iIWdoz77i7KoyLL69eTJn5L0VWL11mxRQy3j/tKovylyEbdX2erCfsh4JJF03uAx09/UkQcAA4A7Ni+OxYuPX/lV5Ym7jxBZCJ6MJqB4TnjvburzCkGkA1FNtz4N7HN7uQ6G87Aws4x11lkm2adrSbsdwMvkHQ58BjwVuBPV/ytDd4lH1sEKkS2APlstuKxhOgFZT/qvZN6prq5Reg8Uf3tT66HqMKr08IbOZSDOLWeBNHRf9eljB32iBhJ+kvga0AOfDwifrhmLZsgKqM6SFMUnPOw2PrUCv8BgmfPz5i9uAp85DFxeypJCzF9OGPm8UDlqTfgap1B9KDssenenFezZScivgJ8ZY3aMrEUgYqAAvKjC0yv8PyQKPMpnn0eKK92J5c4nGEbRAH9Y7DtsQWyotq9D4myN8XchVBk1efyzbbGVhX2VMTJg4SCYipf8dRbSBRTTMZpQ7Oaw96UoMzFcFvGwszK++TDbSKyIDI23e6gdZPD3tDJUzlFv95qr6Dsb66DO9Z9DnsTguiJsi/mdmXMrXT2UFBMBdEPIgvvzttEcNgbCFUXZ5R9sbAD5s8vHGDrHId9HKcFPZsT2YJW/GjeOyGyYZCNAjp6FZZ1l8O+WiVMHc3Y+mSgFS7S6s2XDI4VZAsl6sBNFba5OOxrIBvC1LFyxbBnw3DQJ0W9dxYTcO9FWxx2S1LkUExnv/EGXfbqMyib9FSpw25JKnsw2nJa2Ads6suaHfY1ED0YbslW3DXPF6oDdIxKVLTUOPstoSrYw63ZqWvjBcVAhDbvqVKHfbUE87uC0ZbFt7ctbfAMbH8UerOChWowBH923wAK5s8NiunfXGfFVFAONu9Vjw77agmKrSXF1ibPzSiequ+NHm3SzUcXnM0620Qc9oZUBhqJ/nEY/Wq8D3b9Yz7P3qaT66x3AgZPj7vOqrMtm2GdOewNqAwYQaaS6V+WY49Y0ns26M+WzwXeu/Drx+vst7Ua9shEMTXmMJ8bKOqx8E7e2rrS+fQzKXuCgFBGdHXUng5IdZ1VYycsrdWwj7aIX7x00OYi10zUx3LKfvU1DhUwO8yrf7zYlMeAJkqK62x0/yrCLunjwA3A4Yh4ST1vF/A54DLgYeAtEfGrlV6rHMDxS33OyWy9lGfYljbZsv8L8I/Avy6a917gGxGxvy779F7gb1d6oWxQMHPJsQaLNLNxZIPlN6Yrhj0iviPpstNmvwl4Vf3zrcC3aRD2S7ce5Z+v/ORKTzOzMf3F1qPLPjbuZ/bnRcQTABHxhKQLm/zSNolrpzfx9YhmG2zbGW7sWffkSbpJ0j2S7jly1J/XzTbKuGF/StJFAPX3w8s9MSIORMTeiNh7wXndO+1mtlmMG/YvA2+vf3478KW1aY6ZrZcVwy7pM8B3gSskHZL0DmA/sE/SQ1T12fevbzPNbLWaHI2/cZmHzlx72cwmSqtX0D1dZnz5RGK3Gpm16Oly+Z31VsN+eGE7//Toq9tcpFlSDi88uexjrYZ9WOQ89utz2lykWVKGxfJnvNq9xfV4jv5zZ6uLNEvK8QkJe382uPDeuTYXaZaUR2aXvy2v3S17dPvmf7OJd4Z8+UJ1s0Q47GaJcNjNEuGwmyXCYTdLhMNulgiH3SwRDrtZIhx2s0Q47GaJcNjNEtFkWKpLJH1L0kFJP5R0cz1/l6Q7JT1Ufz93/ZtrZuNqsmUfAe+OiBcB1wDvlPRiTlWFeQHwjXrazCbUimGPiCci4t7652PAQWA3VVWYW+un3Qr88Tq10czWwFl9Zq/LQF0N3MVpVWGAJavCLC4SsTA8scrmmtm4Godd0jbg88C7IuKZpr+3uEjEoD8zThvNbA00CrukPlXQPxURX6hnN64KY2Ybr8nReAEfAw5GxAcWPeSqMGYd0mRYqmuBPwP+V9L99by/o6oCc1tdIeYR4M3r0kIzWxNNKsL8F7BcHVhXhTHrCF9BZ5YIh90sEQ67WSIcdrNEOOxmiXDYzRLhsJslwmE3S4TDbpYIh90sEQ67WSIcdrNEOOxmiXDYzRLhsJslwmE3S4TDbpaIJmPQTUv6b0n/U1eEeX893xVhzDqkyZZ9HnhNRFwJXAVcL+kaXBHGrFOaVISJiDheT/brr8AVYcw6pem48Xk9suxh4M6IcEUYs45pFPaIKCLiKmAP8HJJL2m6AFeEMZsMZ3U0PiKeBr4NXI8rwph1SpOj8RdI2ln/vAV4LfAArghj1ilNKsJcBNwqKad6c7gtIu6Q9F1cEcasM5pUhPk+VZnm0+cfxRVhzDrDV9CZJcJhN0uEw26WCIfdLBEOu1kiHHazRDjsZolw2M0S4bCbJcJhN0uEw26WCIfdLBEOu1kiHHazRDjsZolw2M0S4bCbJaJx2OvhpO+TdEc97YowZh1yNlv2m4GDi6ZdEcasQ5oWidgD/BHw0UWzXRHGrEOabtk/CLwHKBfNc0UYsw5pMm78DcDhiPjeOAtwRRizydBk3PhrgTdKegMwDeyQ9EnqijAR8YQrwphNviZVXG+JiD0RcRnwVuCbEfE2XBHGrFNWc559P7BP0kPAvnrazCZUk93450TEt6kKO7oijFnH+Ao6s0Q47GaJcNjNEuGwmyXCYTdLhMNulgiH3SwRDrtZIhx2s0Q47GaJcNjNEuGwmyXCYTdLhMNulgiH3SwRDrtZIhx2s0Q0GqlG0sPAMaAARhGxV9Iu4HPAZcDDwFsi4lfr00wzW62z2bK/OiKuioi99bQrwph1yGp2410RxqxDmoY9gK9L+p6km+p5rghj1iFNR5e9NiIel3QhcKekB5ouICIOAAcAdmzfHWO00czWQKMte0Q8Xn8/DHwReDl1RRgAV4Qxm3xNar3NSNp+8mfgdcAPcEUYs05pshv/POCLkk4+/9MR8VVJdwO3SXoH8Ajw5vVrppmt1ophj4ifAVcuMd8VYcw6xFfQmSXCYTdLhMNulgiH3SwRDrtZIhx2s0Q47GaJcNjNEuGwmyXCYTdLhMNulgiH3SwRTQevmHhR3ZUHgMJjZJidrtNhj0zVlyBygQQRqAgUoDJQ6eCbQdfDLlEMMhCU/Sr4KoNsGKiEbFhCeEtvBl0Pe0+UU6LMRdkXRV9kRdCbq7bukJEXxUY30xIVEpGLslftdZY9ETmooNr7rDdMWVG20p5Oh320JWP2gpxiIBZ2QLElyOYzpn8Z5PMw9euSfL6oxsY1a5tg4Zwe89sziimY3yWKAfRmoX88yBdg65GCwTPthL3R0XhJOyXdLukBSQclvVLSLkl3Snqo/n7uejf2dGVPDLeK4TYY7ggWdpUMzwmGM2K0pdram22kYiCGM2K4XczvCuYvLJjfFQx3iOG2dv9Hm556+xDw1Yh4IdUQVQdxRRizFVUHj6HsQdkPol9S9qHMq/nR4snvJqPL7gD+EPgYQEQsRMTTuCKM2Yoih+hRBXy6RFsKyumScgDFoDqj1JYm7yvPB44An5B0n6SP1kNKuyKM2Tg26BhSk7D3gJcBH4mIq4ETnMUue0QciIi9EbF30J8Zs5lm3aQCNIJsCNlcRsz2yOcysgXIF2j1OpAmYT8EHIqIu+rp26nC74owZitQVIHPRpAtCC2oCn9RvxG0cyAeaBD2iHgSeFTSFfWs64Af4YowZp3S9Dz7XwGfkjQAfgb8OdUbhSvCmHVEo7BHxP3A3iUeckUYs47wLa5miXDYzRLhsJslwmE3S4TDbpYIh91svWXVDS+RB9GLU9NZdaNMi80ws/XyXKjrG2LIo7rjrZ6Pw96MyvoyxGF1KWI2J7KFenoUrV6KaLac0KkvMiCL6ueWh1vo9Eg1vdmCbU9U9wYXT1UDAaiA/rMFKqD3bOlRamxDPbdVzyGmSjRVUA4yog/lqN1bXLsd9rmC3pzHmLPJFKpHP67DTh70+gXDXo8yD7J6ZOS2dDrsZpNMEeTzQf84qBDF0R6j2ZzBMTF4RuRzkA/b2/V02M3WS8DgWEE+zCh7YvpoNcJsPoR8viQbBb0T7e2ZOuxm6ygbBcyXZEOhQkRW3cte1TaI6vGWOOxm60QRaBRkZUAmNFJ1BD4gK6qQtzlSjcNuto6yooSiOlh3+nnutisVOexmLZiEEmRNhpK+QtL9i76ekfSuSSgSYWbNNRmD7sGIuCoirgJ+F5gFvoiLRJh1ytleLnsd8NOI+D9cJMKsU8427G8FPlP/3KhIhJlNhsZhr0eWfSPw72ezAFeEMZsMZ7Nlfz1wb0Q8VU83KhLhijBmk+Fswn4jp3bhwUUizDqlaX32rcA+4AuLZu8H9kl6qH5s/9o3z8zWStMiEbPAeafNO4qLRJh1RqdHqjGz5hx2s0Q47GaJcNjNEuGwmyXCYTdLhMNulgiH3SwRDrtZIhx2s0Q47GaJaHXAyWKQcWzPVJuLNEtK8f3lt9+thn20BY6+tOXSlWYJGX1r+cfaHUo6D0Y7XYjRbN3kyw9Z3WrYd87M8ie/d0+bizRLyqdnZpd9rNWw7+mfYP/v3N3mIs2S8p3+8uM8thp2IfrK21ykWVLE8sfEmg5L9TeSfijpB5I+I2naFWHMuqVJ+afdwF8DeyPiJUBONX68K8KYdUjTi2p6wBZJPWAr8DiuCGPWKU1qvT0G/APwCPAE8OuI+DquCGPWKU1248+l2opfDlwMzEh6W9MFLK4Ic+Soz7GbbZQmu/GvBX4eEUciYkg1dvzvM0ZFmAvO85F4s43SJOyPANdI2ipJVGPFH8QVYcw6ZcXz7BFxl6TbgXuBEXAfcADYBtwm6R1UbwhvXs+GmtnqNK0I8z7gfafNnscVYcw6w/ezmyXCYTdLhMNulgiH3SwRilj+Zvc1X5h0BDgB/KK1ha6/83F/Jtlm6k+TvlwaERcs9UCrYQeQdE9E7G11oevI/Zlsm6k/q+2Ld+PNEuGwmyViI8J+YAOWuZ7cn8m2mfqzqr60/pndzDaGd+PNEtFq2CVdL+lBST+R1KlhrCRdIulbkg7W4/HdXM/v9Fh8knJJ90m6o57ubH8k7ZR0u6QH6vX0yo73Z03Hfmwt7JJy4MPA64EXAzdKenFby18DI+DdEfEi4BrgnXX7uz4W381Utyyf1OX+fAj4akS8ELiSql+d7M+6jP0YEa18Aa8EvrZo+hbglraWvw79+RKwD3gQuKiedxHw4Ea37Sz6sKf+h3kNcEc9r5P9AXYAP6c+DrVoflf7sxt4FNhFdXfqHcDrVtOfNnfjTzb+pEP1vM6RdBlwNXAX3R6L74PAe4By0byu9uf5wBHgE/XHko9KmqGj/Yl1GPuxzbAvNXp9504FSNoGfB54V0Q8s9HtGZekG4DDEfG9jW7LGukBLwM+EhFXU12W3Yld9qWsduzHpbQZ9kPAJYum91ANSd0ZkvpUQf9URHyhnt1oLL4JdC3wRkkPA58FXiPpk3S3P4eAQxFxVz19O1X4u9qfVY39uJQ2w3438AJJl0saUB1s+HKLy1+Vevy9jwEHI+IDix7q5Fh8EXFLROyJiMuo1sU3I+JtdLc/TwKPSrqinnUd8CM62h/WY+zHlg86vAH4MfBT4O83+iDIWbb9D6g+dnwfuL/+egNwHtVBrofq77s2uq1j9O1VnDpA19n+AFcB99Tr6D+Aczven/cDDwA/AP4NmFpNf3wFnVkifAWdWSIcdrNEOOxmiXDYzRLhsJslwmE3S4TDbpYIh90sEf8P6EO/JfwFzD4AAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Create and wrap the environment\n",
    "env = make_atari('PongNoFrameskip-v4') # only use in no frameskip environment\n",
    "env = wrap_deepmind(env, scale = False, frame_stack=True )\n",
    "n_actions = env.action_space.n\n",
    "state_dim = env.observation_space.shape\n",
    "\n",
    "# env.render()\n",
    "test = env.reset()\n",
    "for i in range(100):\n",
    "    test = env.step(env.action_space.sample())[0]\n",
    "\n",
    "plt.imshow(test._force()[...,0])\n",
    "\n",
    "# plt.imshow(env.render(\"rgb_array\"))\n",
    "# env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DRQN(nn.Module):\n",
    "    def __init__(self, in_channels=4, num_actions=5, device = torch.device(\"cpu\")):\n",
    "        \"\"\"\n",
    "        Initialize a deep Q-learning network as described in\n",
    "        Arguments:\n",
    "            in_channels: number of channel of input.\n",
    "                i.e The number of most recent frames stacked together as describe in the paper\n",
    "            num_actions: number of action-value to output, one-to-one correspondence to action in game.\n",
    "            device: cpu or gpu (cuda:0)\n",
    "        \"\"\"\n",
    "        super(DRQN, self).__init__()\n",
    "        self.device = device\n",
    "        self.num_actions = num_actions\n",
    "        \n",
    "        self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)\n",
    "        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)\n",
    "        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)\n",
    "        self.fc4 = nn.Linear(7 * 7 * 64, 512)\n",
    "        self.gru = nn.GRU(512, num_actions, batch_first=True) # input shape (batch, seq, feature)\n",
    "\n",
    "    def forward(self, x, hidden = None, max_seq = 1, batch_size=1):\n",
    "        # DQN input B*C*feature (32 4 84 84)\n",
    "        # DRQN input B*C*feature (32*seq_len 4 84 84)\n",
    "        x = F.relu(self.conv1(x))\n",
    "        x = F.relu(self.conv2(x))\n",
    "        x = F.relu(self.conv3(x))\n",
    "        x = F.relu(self.fc4(x.reshape(x.size(0), -1)))\n",
    "        hidden = self.init_hidden(batch_size) if hidden is None else hidden\n",
    "        # before go to RNN, reshape the input to (barch, seq, feature)\n",
    "        x = x.reshape(batch_size, max_seq, 512)\n",
    "        return self.gru(x, hidden)\n",
    "\n",
    "    def init_hidden(self, batch_size):\n",
    "        # initialize hidden state to 0\n",
    "        return torch.zeros(1, batch_size, self.num_actions, device= self.device, dtype=torch.float)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Recurrent_Memory_Buffer(object):\n",
    "    # memory buffer to store episodic memory\n",
    "    def __init__(self, memory_size=1000, max_seq = 10):\n",
    "        self.buffer = []\n",
    "        self.memory_size = memory_size\n",
    "        self.max_seq = max_seq\n",
    "        self.next_idx = 0\n",
    "        \n",
    "    def push(self, state, action, reward, next_state, done):\n",
    "        data = (state, action, reward, next_state, done)\n",
    "        if len(self.buffer) <= self.memory_size: # buffer not full\n",
    "            self.buffer.append(data)\n",
    "        else: # buffer is full\n",
    "            self.buffer[self.next_idx] = data\n",
    "        self.next_idx = (self.next_idx + 1) % self.memory_size\n",
    "\n",
    "    def sample(self, batch_size):\n",
    "        # sample episodic memory\n",
    "        states, actions, rewards, next_states, dones = [], [], [], [], []\n",
    "        for i in range(batch_size):\n",
    "            finish = random.randint(self.max_seq, self.size() - 1)\n",
    "            begin = finish-self.max_seq\n",
    "            \n",
    "            data = self.rec_memory_buffer.buffer[begin:finish]\n",
    "            state, action, reward, next_state, done= zip(*data)\n",
    "            states.append(np.concatenate([self.observe(state_i) for state_i in state]))\n",
    "            actions.append(action)\n",
    "            rewards.append(reward)\n",
    "            next_states.append(np.concatenate([self.observe(state_i) for state_i in next_state]))\n",
    "            dones.append(done)\n",
    "\n",
    "        return np.concatenate(states), actions, rewards, np.concatenate(next_states), dones\n",
    "    \n",
    "    def size(self):\n",
    "        return len(self.buffer)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DRQNAgent: \n",
    "    # DRQN agent\n",
    "    def __init__(self, in_channels = 1, action_space = [], USE_CUDA = False, memory_size = 10000, epsilon  = 1, lr = 1e-4, max_seq = 10, batch_size = 32):\n",
    "        self.USE_CUDA = USE_CUDA\n",
    "        self.device = torch.device(\"cuda:0\" if USE_CUDA else \"cpu\")\n",
    "        self.max_seq = max_seq\n",
    "        self.batch_size = batch_size\n",
    "        self.epsilon = epsilon\n",
    "        self.action_space = action_space\n",
    "        self.rec_memory_buffer = Recurrent_Memory_Buffer(memory_size)\n",
    "        self.DRQN = DRQN(in_channels = in_channels, num_actions = action_space.n, device = self.device)\n",
    "        self.DRQN_target = DRQN(in_channels = in_channels, num_actions = action_space.n, device = self.device)\n",
    "        self.DRQN_target.load_state_dict(self.DRQN.state_dict())\n",
    "\n",
    "        if USE_CUDA:\n",
    "            self.DRQN = self.DRQN.cuda()\n",
    "            self.DRQN_target = self.DRQN_target.cuda()\n",
    "        self.optimizer = optim.RMSprop(self.DRQN.parameters(),lr=lr, eps=0.001, alpha=0.95)\n",
    "\n",
    "    def observe(self, lazyframe):\n",
    "        # from Lazy frame to tensor\n",
    "        state =  torch.from_numpy(lazyframe._force().transpose(2,0,1)[None]/255).float()\n",
    "        if self.USE_CUDA:\n",
    "            state = state.cuda()\n",
    "        return state\n",
    "\n",
    "    def value(self, state, hidden = None):\n",
    "        # get q_values of a given state\n",
    "        q_values, hidden = self.DRQN(state, hidden)\n",
    "        return q_values, hidden\n",
    "    \n",
    "    def act(self, state, hidden = None, epsilon = None):\n",
    "        \"\"\"\n",
    "        sample actions with epsilon-greedy policy\n",
    "        recap: with p = epsilon pick random action, else pick action with highest Q(s,a)\n",
    "        \"\"\"\n",
    "        if epsilon is None: epsilon = self.epsilon\n",
    "        q_values, hidden  = self.value(state, hidden)\n",
    "        q_values = q_values.squeeze(0).cpu().detach().numpy()\n",
    "        if random.random()<epsilon:\n",
    "            aciton = random.randrange(self.action_space.n)\n",
    "        else:\n",
    "            aciton = q_values.argmax(1)[0]\n",
    "        return aciton, hidden\n",
    "    \n",
    "    def compute_td_loss(self, states, actions, rewards, next_states, is_done, gamma=0.99):\n",
    "        \"\"\" Compute td loss using torch operations only.\"\"\"\n",
    "        actions = torch.tensor(actions).long().flatten()    # shape: [batch_size * seq_len]\n",
    "        rewards = torch.tensor(rewards, dtype =torch.float).flatten()  # shape: [batch_size * seq_len]\n",
    "        is_done = torch.tensor(is_done).bool().flatten()  # shape: [batch_size * seq_len]\n",
    "        \n",
    "        if self.USE_CUDA:\n",
    "            actions = actions.cuda()\n",
    "            rewards = rewards.cuda()\n",
    "            is_done = is_done.cuda()\n",
    "\n",
    "        # get q-values for all actions in current states\n",
    "        predicted_qvalues, hidden = self.DRQN(states, batch_size = self.batch_size, max_seq = self.max_seq)\n",
    "        predicted_qvalues = predicted_qvalues.reshape(-1, self.action_space.n)\n",
    "        predicted_qvalues = predicted_qvalues.squeeze(0)\n",
    "        \n",
    "        # select q-values for chosen actions\n",
    "        predicted_qvalues_for_actions = predicted_qvalues[\n",
    "          range(states.shape[0]), actions\n",
    "        ]\n",
    "\n",
    "        # compute q-values for all actions in next states\n",
    "        predicted_next_qvalues, hidden_target = self.DRQN_target(next_states, batch_size = self.batch_size, max_seq = self.max_seq) # YOUR CODE\n",
    "        predicted_next_qvalues = predicted_next_qvalues.squeeze(0)\n",
    "        predicted_next_qvalues = predicted_next_qvalues.reshape(-1, self.action_space.n)\n",
    "\n",
    "        # compute V*(next_states) using predicted next q-values\n",
    "        next_state_values =  predicted_next_qvalues.max(-1)[0] \n",
    "\n",
    "        # compute \"target q-values\" for loss - it's what's inside square parentheses in the above formula.\n",
    "        target_qvalues_for_actions = rewards + gamma *next_state_values \n",
    "\n",
    "        # at the last state we shall use simplified formula: Q(s,a) = r(s,a) since s' doesn't exist\n",
    "        target_qvalues_for_actions = torch.where(\n",
    "            is_done, rewards, target_qvalues_for_actions)\n",
    "\n",
    "        # mean squared error loss to minimize\n",
    "        loss = F.smooth_l1_loss(predicted_qvalues_for_actions, target_qvalues_for_actions.detach())\n",
    "\n",
    "        return loss\n",
    "    \n",
    "    def sample_from_buffer(self, batch_size):\n",
    "        # rewriten sample() in buffer with pytorch operations\n",
    "        states, actions, rewards, next_states, dones = [], [], [], [], []\n",
    "        \n",
    "        for i in range(batch_size):\n",
    "            finish = random.randint(self.max_seq, self.rec_memory_buffer.size() - 1)\n",
    "            begin = finish-self.max_seq\n",
    "            \n",
    "            data = self.rec_memory_buffer.buffer[begin:finish]\n",
    "            state, action, reward, next_state, done= zip(*data)\n",
    "            states.append(torch.cat([self.observe(state_i) for state_i in state]))\n",
    "            actions.append(action)\n",
    "            rewards.append(reward)\n",
    "            next_states.append(torch.cat([self.observe(state_i) for state_i in next_state]))\n",
    "            dones.append(done)\n",
    "\n",
    "        return torch.cat(states), actions, rewards, torch.cat(next_states), dones\n",
    "\n",
    "    def learn_from_experience(self, batch_size):\n",
    "        # learn from experience\n",
    "        if self.rec_memory_buffer.size() > batch_size:\n",
    "            states, actions, rewards, next_states, dones = self.sample_from_buffer(batch_size)\n",
    "            td_loss = self.compute_td_loss(states, actions, rewards, next_states, dones)\n",
    "            self.optimizer.zero_grad()\n",
    "            td_loss.backward()\n",
    "            for param in self.DRQN.parameters():\n",
    "                # clip the gradient\n",
    "                param.grad.data.clamp_(-1, 1)\n",
    "            self.optimizer.step()\n",
    "            return(td_loss.item())\n",
    "        else:\n",
    "            return(0)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:NaN or Inf found in input tensor.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "frames:     0, reward:   nan, loss: 0.000000, epsilon: 1.000000, episode:    0\n",
      "frames:  1000, reward: -21.000000, loss: 0.000000, epsilon: 0.967544, episode:    1\n",
      "frames:  2000, reward: -20.000000, loss: 0.000000, epsilon: 0.936152, episode:    2\n",
      "frames:  3000, reward: -20.000000, loss: 0.000000, epsilon: 0.905789, episode:    3\n",
      "frames:  4000, reward: -20.250000, loss: 0.000000, epsilon: 0.876422, episode:    4\n",
      "frames:  5000, reward: -20.000000, loss: 0.021156, epsilon: 0.848017, episode:    5\n",
      "frames:  6000, reward: -19.833333, loss: 0.011389, epsilon: 0.820543, episode:    6\n",
      "frames:  7000, reward: -20.000000, loss: 0.007602, epsilon: 0.793971, episode:    7\n",
      "frames:  8000, reward: -20.000000, loss: 0.007895, epsilon: 0.768269, episode:    8\n",
      "frames:  9000, reward: -20.200000, loss: 0.007752, epsilon: 0.743410, episode:   10\n",
      "frames: 10000, reward: -20.200000, loss: 0.009361, epsilon: 0.719366, episode:   11\n",
      "frames: 11000, reward: -20.300000, loss: 0.011069, epsilon: 0.696110, episode:   12\n",
      "frames: 12000, reward: -20.400000, loss: 0.011838, epsilon: 0.673617, episode:   13\n",
      "frames: 13000, reward: -20.300000, loss: 0.010123, epsilon: 0.651861, episode:   14\n",
      "frames: 14000, reward: -20.400000, loss: 0.003541, epsilon: 0.630818, episode:   15\n",
      "frames: 15000, reward: -20.600000, loss: 0.007865, epsilon: 0.610465, episode:   16\n",
      "frames: 16000, reward: -20.400000, loss: 0.005524, epsilon: 0.590780, episode:   17\n",
      "frames: 17000, reward: -20.500000, loss: 0.006949, epsilon: 0.571740, episode:   19\n",
      "frames: 18000, reward: -20.500000, loss: 0.006297, epsilon: 0.553324, episode:   20\n",
      "frames: 19000, reward: -20.500000, loss: 0.005883, epsilon: 0.535511, episode:   21\n",
      "frames: 20000, reward: -20.600000, loss: 0.002443, epsilon: 0.518283, episode:   22\n",
      "frames: 21000, reward: -20.500000, loss: 0.002456, epsilon: 0.501619, episode:   23\n",
      "frames: 22000, reward: -20.600000, loss: 0.005527, epsilon: 0.485502, episode:   24\n",
      "frames: 23000, reward: -20.700000, loss: 0.004897, epsilon: 0.469913, episode:   25\n",
      "frames: 24000, reward: -20.900000, loss: 0.007468, epsilon: 0.454836, episode:   27\n",
      "frames: 25000, reward: -20.900000, loss: 0.004536, epsilon: 0.440252, episode:   28\n",
      "frames: 26000, reward: -20.800000, loss: 0.001681, epsilon: 0.426147, episode:   29\n",
      "frames: 27000, reward: -20.600000, loss: 0.001582, epsilon: 0.412504, episode:   30\n",
      "frames: 28000, reward: -20.500000, loss: 0.003278, epsilon: 0.399308, episode:   31\n",
      "frames: 29000, reward: -20.500000, loss: 0.001133, epsilon: 0.386545, episode:   32\n",
      "frames: 30000, reward: -20.400000, loss: 0.003491, epsilon: 0.374201, episode:   33\n",
      "frames: 31000, reward: -20.400000, loss: 0.000908, epsilon: 0.362261, episode:   34\n",
      "frames: 32000, reward: -20.300000, loss: 0.000776, epsilon: 0.350712, episode:   36\n",
      "frames: 33000, reward: -20.300000, loss: 0.001200, epsilon: 0.339542, episode:   37\n",
      "frames: 34000, reward: -20.100000, loss: 0.001053, epsilon: 0.328739, episode:   38\n",
      "frames: 35000, reward: -20.100000, loss: 0.001659, epsilon: 0.318289, episode:   39\n",
      "frames: 36000, reward: -20.000000, loss: 0.002701, epsilon: 0.308182, episode:   40\n",
      "frames: 37000, reward: -20.100000, loss: 0.001461, epsilon: 0.298407, episode:   41\n",
      "frames: 38000, reward: -20.000000, loss: 0.001208, epsilon: 0.288952, episode:   42\n",
      "frames: 39000, reward: -20.100000, loss: 0.002670, epsilon: 0.279806, episode:   43\n",
      "frames: 40000, reward: -19.600000, loss: 0.004548, epsilon: 0.270961, episode:   44\n",
      "frames: 41000, reward: -19.700000, loss: 0.002634, epsilon: 0.262406, episode:   45\n",
      "frames: 42000, reward: -19.600000, loss: 0.002597, epsilon: 0.254131, episode:   46\n",
      "frames: 43000, reward: -19.500000, loss: 0.003329, epsilon: 0.246127, episode:   47\n",
      "frames: 44000, reward: -19.700000, loss: 0.001314, epsilon: 0.238386, episode:   48\n",
      "frames: 45000, reward: -19.600000, loss: 0.002580, epsilon: 0.230899, episode:   49\n",
      "frames: 46000, reward: -19.800000, loss: 0.001514, epsilon: 0.223657, episode:   50\n",
      "frames: 47000, reward: -19.800000, loss: 0.001016, epsilon: 0.216652, episode:   50\n",
      "frames: 48000, reward: -19.600000, loss: 0.005111, epsilon: 0.209878, episode:   51\n",
      "frames: 49000, reward: -19.400000, loss: 0.002114, epsilon: 0.203325, episode:   52\n",
      "frames: 50000, reward: -19.200000, loss: 0.002325, epsilon: 0.196987, episode:   53\n",
      "frames: 51000, reward: -19.200000, loss: 0.001649, epsilon: 0.190857, episode:   53\n",
      "frames: 52000, reward: -19.100000, loss: 0.002247, epsilon: 0.184928, episode:   54\n",
      "frames: 53000, reward: -18.900000, loss: 0.001891, epsilon: 0.179193, episode:   55\n",
      "frames: 54000, reward: -18.900000, loss: 0.002428, epsilon: 0.173646, episode:   56\n",
      "frames: 55000, reward: -18.900000, loss: 0.003432, epsilon: 0.168281, episode:   56\n",
      "frames: 56000, reward: -18.600000, loss: 0.001659, epsilon: 0.163092, episode:   57\n",
      "frames: 57000, reward: -18.600000, loss: 0.002859, epsilon: 0.158073, episode:   57\n",
      "frames: 58000, reward: -18.000000, loss: 0.002342, epsilon: 0.153219, episode:   58\n",
      "frames: 59000, reward: -17.800000, loss: 0.006105, epsilon: 0.148523, episode:   59\n",
      "frames: 60000, reward: -17.700000, loss: 0.001271, epsilon: 0.143982, episode:   60\n",
      "frames: 61000, reward: -17.700000, loss: 0.002508, epsilon: 0.139589, episode:   60\n",
      "frames: 62000, reward: -17.600000, loss: 0.003631, epsilon: 0.135341, episode:   61\n",
      "frames: 63000, reward: -17.600000, loss: 0.001515, epsilon: 0.131232, episode:   62\n",
      "frames: 64000, reward: -17.700000, loss: 0.003296, epsilon: 0.127257, episode:   63\n",
      "frames: 65000, reward: -17.700000, loss: 0.001446, epsilon: 0.123413, episode:   63\n",
      "frames: 66000, reward: -18.100000, loss: 0.001449, epsilon: 0.119695, episode:   64\n",
      "frames: 67000, reward: -18.000000, loss: 0.001999, epsilon: 0.116099, episode:   65\n",
      "frames: 68000, reward: -18.000000, loss: 0.002700, epsilon: 0.112621, episode:   65\n",
      "frames: 69000, reward: -17.300000, loss: 0.002082, epsilon: 0.109256, episode:   66\n",
      "frames: 70000, reward: -17.300000, loss: 0.003149, epsilon: 0.106002, episode:   66\n",
      "frames: 71000, reward: -17.500000, loss: 0.001861, epsilon: 0.102855, episode:   67\n",
      "frames: 72000, reward: -17.500000, loss: 0.001065, epsilon: 0.099811, episode:   67\n",
      "frames: 73000, reward: -17.800000, loss: 0.001042, epsilon: 0.096866, episode:   68\n",
      "frames: 74000, reward: -17.800000, loss: 0.002277, epsilon: 0.094019, episode:   68\n",
      "frames: 75000, reward: -17.900000, loss: 0.000738, epsilon: 0.091264, episode:   69\n",
      "frames: 76000, reward: -17.800000, loss: 0.000693, epsilon: 0.088600, episode:   70\n",
      "frames: 77000, reward: -17.800000, loss: 0.002340, epsilon: 0.086023, episode:   70\n",
      "frames: 78000, reward: -17.800000, loss: 0.000897, epsilon: 0.083531, episode:   71\n",
      "frames: 79000, reward: -17.800000, loss: 0.002513, epsilon: 0.081120, episode:   71\n",
      "frames: 80000, reward: -17.700000, loss: 0.002786, epsilon: 0.078789, episode:   72\n",
      "frames: 81000, reward: -17.700000, loss: 0.000622, epsilon: 0.076533, episode:   73\n",
      "frames: 82000, reward: -17.700000, loss: 0.001471, epsilon: 0.074352, episode:   73\n",
      "frames: 83000, reward: -17.600000, loss: 0.000910, epsilon: 0.072243, episode:   74\n",
      "frames: 84000, reward: -17.600000, loss: 0.001152, epsilon: 0.070202, episode:   74\n",
      "frames: 85000, reward: -17.600000, loss: 0.002030, epsilon: 0.068228, episode:   75\n",
      "frames: 86000, reward: -17.600000, loss: 0.001790, epsilon: 0.066319, episode:   75\n",
      "frames: 87000, reward: -17.900000, loss: 0.000769, epsilon: 0.064473, episode:   76\n",
      "frames: 88000, reward: -17.700000, loss: 0.001061, epsilon: 0.062687, episode:   77\n",
      "frames: 89000, reward: -17.700000, loss: 0.002388, epsilon: 0.060960, episode:   77\n",
      "frames: 90000, reward: -17.500000, loss: 0.000624, epsilon: 0.059289, episode:   78\n",
      "frames: 91000, reward: -17.500000, loss: 0.000917, epsilon: 0.057673, episode:   78\n",
      "frames: 92000, reward: -17.500000, loss: 0.000747, epsilon: 0.056110, episode:   79\n",
      "frames: 93000, reward: -17.500000, loss: 0.002093, epsilon: 0.054599, episode:   80\n",
      "frames: 94000, reward: -17.500000, loss: 0.002966, epsilon: 0.053137, episode:   80\n",
      "frames: 95000, reward: -17.100000, loss: 0.002406, epsilon: 0.051722, episode:   81\n",
      "frames: 96000, reward: -17.100000, loss: 0.000736, epsilon: 0.050355, episode:   81\n",
      "frames: 97000, reward: -16.500000, loss: 0.000631, epsilon: 0.049032, episode:   82\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "frames: 98000, reward: -16.500000, loss: 0.003315, epsilon: 0.047752, episode:   82\n",
      "frames: 99000, reward: -16.500000, loss: 0.003370, epsilon: 0.046514, episode:   82\n",
      "frames: 100000, reward: -15.300000, loss: 0.002055, epsilon: 0.045317, episode:   83\n",
      "frames: 101000, reward: -15.300000, loss: 0.002854, epsilon: 0.044159, episode:   83\n",
      "frames: 102000, reward: -15.000000, loss: 0.000711, epsilon: 0.043040, episode:   84\n",
      "frames: 103000, reward: -14.700000, loss: 0.000445, epsilon: 0.041956, episode:   85\n",
      "frames: 104000, reward: -14.700000, loss: 0.001948, epsilon: 0.040909, episode:   85\n",
      "frames: 105000, reward: -15.000000, loss: 0.001191, epsilon: 0.039895, episode:   86\n",
      "frames: 106000, reward: -15.000000, loss: 0.003230, epsilon: 0.038915, episode:   86\n",
      "frames: 107000, reward: -15.200000, loss: 0.000588, epsilon: 0.037967, episode:   87\n",
      "frames: 108000, reward: -15.200000, loss: 0.000731, epsilon: 0.037050, episode:   87\n",
      "frames: 109000, reward: -14.800000, loss: 0.001891, epsilon: 0.036164, episode:   88\n",
      "frames: 110000, reward: -14.800000, loss: 0.003577, epsilon: 0.035306, episode:   88\n",
      "frames: 111000, reward: -14.500000, loss: 0.002055, epsilon: 0.034476, episode:   89\n",
      "frames: 112000, reward: -14.700000, loss: 0.000439, epsilon: 0.033674, episode:   90\n",
      "frames: 113000, reward: -14.700000, loss: 0.000611, epsilon: 0.032898, episode:   90\n",
      "frames: 114000, reward: -14.700000, loss: 0.002191, epsilon: 0.032147, episode:   90\n",
      "frames: 115000, reward: -14.400000, loss: 0.001571, epsilon: 0.031421, episode:   91\n",
      "frames: 116000, reward: -14.400000, loss: 0.001685, epsilon: 0.030719, episode:   91\n",
      "frames: 117000, reward: -15.100000, loss: 0.001438, epsilon: 0.030039, episode:   92\n",
      "frames: 118000, reward: -16.200000, loss: 0.000744, epsilon: 0.029383, episode:   93\n",
      "frames: 119000, reward: -16.200000, loss: 0.000571, epsilon: 0.028747, episode:   93\n",
      "frames: 120000, reward: -16.200000, loss: 0.000325, epsilon: 0.028132, episode:   93\n",
      "frames: 121000, reward: -15.800000, loss: 0.001684, epsilon: 0.027538, episode:   94\n",
      "frames: 122000, reward: -15.800000, loss: 0.003403, epsilon: 0.026963, episode:   94\n",
      "frames: 123000, reward: -15.800000, loss: 0.001908, epsilon: 0.026407, episode:   94\n",
      "frames: 124000, reward: -15.800000, loss: 0.000278, epsilon: 0.025869, episode:   95\n",
      "frames: 125000, reward: -15.800000, loss: 0.000563, epsilon: 0.025349, episode:   95\n",
      "frames: 126000, reward: -15.400000, loss: 0.000390, epsilon: 0.024846, episode:   96\n",
      "frames: 127000, reward: -15.400000, loss: 0.000504, epsilon: 0.024359, episode:   96\n",
      "frames: 128000, reward: -15.000000, loss: 0.000416, epsilon: 0.023888, episode:   97\n",
      "frames: 129000, reward: -15.000000, loss: 0.000910, epsilon: 0.023433, episode:   97\n",
      "frames: 130000, reward: -15.600000, loss: 0.002919, epsilon: 0.022992, episode:   98\n",
      "frames: 131000, reward: -15.600000, loss: 0.002191, epsilon: 0.022567, episode:   98\n",
      "frames: 132000, reward: -15.700000, loss: 0.001661, epsilon: 0.022155, episode:   99\n",
      "frames: 133000, reward: -15.700000, loss: 0.001477, epsilon: 0.021756, episode:   99\n",
      "frames: 134000, reward: -15.700000, loss: 0.000880, epsilon: 0.021371, episode:   99\n",
      "frames: 135000, reward: -15.300000, loss: 0.000512, epsilon: 0.020998, episode:  100\n",
      "frames: 136000, reward: -16.000000, loss: 0.001365, epsilon: 0.020637, episode:  101\n",
      "frames: 137000, reward: -16.000000, loss: 0.001843, epsilon: 0.020289, episode:  101\n",
      "frames: 138000, reward: -16.100000, loss: 0.000807, epsilon: 0.019951, episode:  102\n",
      "frames: 139000, reward: -16.100000, loss: 0.000692, epsilon: 0.019625, episode:  102\n",
      "frames: 140000, reward: -16.100000, loss: 0.000428, epsilon: 0.019310, episode:  103\n",
      "frames: 141000, reward: -16.100000, loss: 0.001678, epsilon: 0.019004, episode:  103\n",
      "frames: 142000, reward: -16.100000, loss: 0.000566, epsilon: 0.018709, episode:  104\n",
      "frames: 143000, reward: -16.100000, loss: 0.004263, epsilon: 0.018424, episode:  104\n",
      "frames: 144000, reward: -16.300000, loss: 0.000338, epsilon: 0.018147, episode:  105\n",
      "frames: 145000, reward: -16.300000, loss: 0.000300, epsilon: 0.017880, episode:  105\n",
      "frames: 146000, reward: -16.300000, loss: 0.001338, epsilon: 0.017622, episode:  105\n",
      "frames: 147000, reward: -16.000000, loss: 0.000352, epsilon: 0.017372, episode:  106\n",
      "frames: 148000, reward: -16.000000, loss: 0.000256, epsilon: 0.017130, episode:  106\n",
      "frames: 149000, reward: -16.100000, loss: 0.004552, epsilon: 0.016897, episode:  107\n",
      "frames: 150000, reward: -16.100000, loss: 0.000399, epsilon: 0.016671, episode:  107\n",
      "frames: 151000, reward: -15.800000, loss: 0.001671, epsilon: 0.016452, episode:  108\n",
      "frames: 152000, reward: -15.800000, loss: 0.002406, epsilon: 0.016240, episode:  108\n",
      "frames: 153000, reward: -15.800000, loss: 0.000368, epsilon: 0.016036, episode:  108\n",
      "frames: 154000, reward: -15.400000, loss: 0.001936, epsilon: 0.015838, episode:  109\n",
      "frames: 155000, reward: -15.400000, loss: 0.000327, epsilon: 0.015647, episode:  109\n",
      "frames: 156000, reward: -15.400000, loss: 0.000344, epsilon: 0.015461, episode:  109\n",
      "frames: 157000, reward: -15.300000, loss: 0.001470, epsilon: 0.015282, episode:  110\n",
      "frames: 158000, reward: -15.300000, loss: 0.001326, epsilon: 0.015109, episode:  110\n",
      "frames: 159000, reward: -14.400000, loss: 0.002250, epsilon: 0.014942, episode:  111\n",
      "frames: 160000, reward: -14.400000, loss: 0.000334, epsilon: 0.014780, episode:  111\n",
      "frames: 161000, reward: -14.400000, loss: 0.001744, epsilon: 0.014623, episode:  111\n",
      "frames: 162000, reward: -14.400000, loss: 0.000511, epsilon: 0.014471, episode:  111\n",
      "frames: 163000, reward: -13.300000, loss: 0.000414, epsilon: 0.014325, episode:  112\n",
      "frames: 164000, reward: -13.300000, loss: 0.002788, epsilon: 0.014183, episode:  112\n",
      "frames: 165000, reward: -12.900000, loss: 0.000322, epsilon: 0.014046, episode:  113\n",
      "frames: 166000, reward: -12.900000, loss: 0.001817, epsilon: 0.013913, episode:  113\n",
      "frames: 167000, reward: -12.900000, loss: 0.001485, epsilon: 0.013785, episode:  113\n",
      "frames: 168000, reward: -13.200000, loss: 0.000242, epsilon: 0.013661, episode:  114\n",
      "frames: 169000, reward: -13.200000, loss: 0.000469, epsilon: 0.013541, episode:  114\n",
      "frames: 170000, reward: -12.600000, loss: 0.000515, epsilon: 0.013425, episode:  115\n",
      "frames: 171000, reward: -12.600000, loss: 0.001588, epsilon: 0.013313, episode:  115\n",
      "frames: 172000, reward: -12.600000, loss: 0.000430, epsilon: 0.013204, episode:  115\n",
      "frames: 173000, reward: -12.400000, loss: 0.001576, epsilon: 0.013099, episode:  116\n",
      "frames: 174000, reward: -12.400000, loss: 0.001631, epsilon: 0.012997, episode:  116\n",
      "frames: 175000, reward: -12.400000, loss: 0.000217, epsilon: 0.012899, episode:  116\n",
      "frames: 176000, reward: -12.000000, loss: 0.003076, epsilon: 0.012804, episode:  117\n",
      "frames: 177000, reward: -12.000000, loss: 0.001716, epsilon: 0.012712, episode:  117\n",
      "frames: 178000, reward: -11.900000, loss: 0.002436, epsilon: 0.012623, episode:  118\n",
      "frames: 179000, reward: -11.900000, loss: 0.002400, epsilon: 0.012537, episode:  118\n",
      "frames: 180000, reward: -12.000000, loss: 0.000308, epsilon: 0.012454, episode:  119\n",
      "frames: 181000, reward: -12.000000, loss: 0.001421, epsilon: 0.012374, episode:  119\n",
      "frames: 182000, reward: -12.000000, loss: 0.000433, epsilon: 0.012296, episode:  119\n",
      "frames: 183000, reward: -12.000000, loss: 0.001936, epsilon: 0.012220, episode:  119\n",
      "frames: 184000, reward: -11.400000, loss: 0.004054, epsilon: 0.012148, episode:  120\n",
      "frames: 185000, reward: -11.400000, loss: 0.001605, epsilon: 0.012077, episode:  120\n",
      "frames: 186000, reward: -12.400000, loss: 0.004000, epsilon: 0.012009, episode:  121\n",
      "frames: 187000, reward: -12.400000, loss: 0.003102, epsilon: 0.011943, episode:  121\n",
      "frames: 188000, reward: -13.100000, loss: 0.000440, epsilon: 0.011880, episode:  122\n",
      "frames: 189000, reward: -13.100000, loss: 0.000530, epsilon: 0.011818, episode:  122\n",
      "frames: 190000, reward: -13.100000, loss: 0.001865, epsilon: 0.011758, episode:  122\n",
      "frames: 191000, reward: -12.600000, loss: 0.001245, epsilon: 0.011701, episode:  123\n",
      "frames: 192000, reward: -12.600000, loss: 0.002741, epsilon: 0.011645, episode:  123\n",
      "frames: 193000, reward: -12.600000, loss: 0.003426, epsilon: 0.011591, episode:  123\n",
      "frames: 194000, reward: -12.300000, loss: 0.005237, epsilon: 0.011539, episode:  124\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "frames: 195000, reward: -12.300000, loss: 0.000304, epsilon: 0.011488, episode:  124\n",
      "frames: 196000, reward: -12.300000, loss: 0.001699, epsilon: 0.011440, episode:  124\n",
      "frames: 197000, reward: -12.300000, loss: 0.000275, epsilon: 0.011392, episode:  125\n",
      "frames: 198000, reward: -12.300000, loss: 0.001613, epsilon: 0.011347, episode:  125\n",
      "frames: 199000, reward: -12.300000, loss: 0.001569, epsilon: 0.011303, episode:  125\n",
      "frames: 200000, reward: -12.700000, loss: 0.001021, epsilon: 0.011260, episode:  126\n",
      "frames: 201000, reward: -12.700000, loss: 0.000621, epsilon: 0.011219, episode:  126\n",
      "frames: 202000, reward: -12.900000, loss: 0.001756, epsilon: 0.011179, episode:  127\n",
      "frames: 203000, reward: -12.900000, loss: 0.001595, epsilon: 0.011140, episode:  127\n",
      "frames: 204000, reward: -12.900000, loss: 0.000179, epsilon: 0.011103, episode:  127\n",
      "frames: 205000, reward: -12.900000, loss: 0.000375, epsilon: 0.011066, episode:  128\n",
      "frames: 206000, reward: -12.900000, loss: 0.001476, epsilon: 0.011032, episode:  128\n",
      "frames: 207000, reward: -12.900000, loss: 0.003844, epsilon: 0.010998, episode:  128\n",
      "frames: 208000, reward: -12.600000, loss: 0.000238, epsilon: 0.010965, episode:  129\n",
      "frames: 209000, reward: -12.600000, loss: 0.000340, epsilon: 0.010933, episode:  129\n",
      "frames: 210000, reward: -13.100000, loss: 0.000363, epsilon: 0.010903, episode:  130\n",
      "frames: 211000, reward: -13.100000, loss: 0.000240, epsilon: 0.010873, episode:  130\n",
      "frames: 212000, reward: -13.100000, loss: 0.000274, epsilon: 0.010845, episode:  130\n",
      "frames: 213000, reward: -12.700000, loss: 0.001585, epsilon: 0.010817, episode:  131\n",
      "frames: 214000, reward: -12.700000, loss: 0.000256, epsilon: 0.010790, episode:  131\n",
      "frames: 215000, reward: -12.700000, loss: 0.004126, epsilon: 0.010764, episode:  131\n",
      "frames: 216000, reward: -12.000000, loss: 0.002103, epsilon: 0.010739, episode:  132\n",
      "frames: 217000, reward: -12.000000, loss: 0.003460, epsilon: 0.010715, episode:  132\n",
      "frames: 218000, reward: -12.800000, loss: 0.003094, epsilon: 0.010691, episode:  133\n",
      "frames: 219000, reward: -12.800000, loss: 0.001469, epsilon: 0.010669, episode:  133\n",
      "frames: 220000, reward: -12.800000, loss: 0.001812, epsilon: 0.010647, episode:  133\n",
      "frames: 221000, reward: -12.000000, loss: 0.001393, epsilon: 0.010626, episode:  134\n",
      "frames: 222000, reward: -12.000000, loss: 0.001613, epsilon: 0.010605, episode:  134\n",
      "frames: 223000, reward: -11.900000, loss: 0.002569, epsilon: 0.010585, episode:  135\n",
      "frames: 224000, reward: -11.900000, loss: 0.000339, epsilon: 0.010566, episode:  135\n",
      "frames: 225000, reward: -11.900000, loss: 0.004651, epsilon: 0.010548, episode:  135\n",
      "frames: 226000, reward: -11.900000, loss: 0.003060, epsilon: 0.010530, episode:  136\n",
      "frames: 227000, reward: -11.900000, loss: 0.003145, epsilon: 0.010512, episode:  136\n",
      "frames: 228000, reward: -11.900000, loss: 0.000252, epsilon: 0.010495, episode:  136\n",
      "frames: 229000, reward: -11.500000, loss: 0.001803, epsilon: 0.010479, episode:  137\n",
      "frames: 230000, reward: -11.500000, loss: 0.001502, epsilon: 0.010463, episode:  137\n",
      "frames: 231000, reward: -11.600000, loss: 0.000633, epsilon: 0.010448, episode:  138\n",
      "frames: 232000, reward: -11.600000, loss: 0.000435, epsilon: 0.010434, episode:  138\n",
      "frames: 233000, reward: -11.600000, loss: 0.001938, epsilon: 0.010419, episode:  138\n",
      "frames: 234000, reward: -12.000000, loss: 0.001459, epsilon: 0.010406, episode:  139\n",
      "frames: 235000, reward: -12.000000, loss: 0.003660, epsilon: 0.010392, episode:  139\n",
      "frames: 236000, reward: -12.000000, loss: 0.000450, epsilon: 0.010379, episode:  140\n",
      "frames: 237000, reward: -12.000000, loss: 0.000380, epsilon: 0.010367, episode:  140\n",
      "frames: 238000, reward: -12.000000, loss: 0.004056, epsilon: 0.010355, episode:  140\n",
      "frames: 239000, reward: -12.000000, loss: 0.001433, epsilon: 0.010343, episode:  141\n",
      "frames: 240000, reward: -12.000000, loss: 0.001430, epsilon: 0.010332, episode:  141\n",
      "frames: 241000, reward: -12.700000, loss: 0.004516, epsilon: 0.010321, episode:  142\n",
      "frames: 242000, reward: -12.700000, loss: 0.000212, epsilon: 0.010311, episode:  142\n",
      "frames: 243000, reward: -12.700000, loss: 0.004136, epsilon: 0.010301, episode:  142\n",
      "frames: 244000, reward: -11.900000, loss: 0.000233, epsilon: 0.010291, episode:  143\n",
      "frames: 245000, reward: -11.900000, loss: 0.002570, epsilon: 0.010281, episode:  143\n",
      "frames: 246000, reward: -11.900000, loss: 0.001560, epsilon: 0.010272, episode:  143\n",
      "frames: 247000, reward: -12.700000, loss: 0.000200, epsilon: 0.010263, episode:  144\n",
      "frames: 248000, reward: -12.700000, loss: 0.001321, epsilon: 0.010254, episode:  144\n",
      "frames: 249000, reward: -12.800000, loss: 0.002350, epsilon: 0.010246, episode:  145\n",
      "frames: 250000, reward: -12.800000, loss: 0.001277, epsilon: 0.010238, episode:  145\n",
      "frames: 251000, reward: -12.800000, loss: 0.000239, epsilon: 0.010230, episode:  145\n",
      "frames: 252000, reward: -12.200000, loss: 0.002547, epsilon: 0.010223, episode:  146\n",
      "frames: 253000, reward: -12.200000, loss: 0.001515, epsilon: 0.010215, episode:  146\n",
      "frames: 254000, reward: -12.200000, loss: 0.001339, epsilon: 0.010208, episode:  146\n",
      "frames: 255000, reward: -12.100000, loss: 0.001516, epsilon: 0.010201, episode:  147\n",
      "frames: 256000, reward: -12.100000, loss: 0.002016, epsilon: 0.010195, episode:  147\n",
      "frames: 257000, reward: -12.100000, loss: 0.002609, epsilon: 0.010188, episode:  147\n",
      "frames: 258000, reward: -12.100000, loss: 0.000284, epsilon: 0.010182, episode:  147\n",
      "frames: 259000, reward: -10.900000, loss: 0.002670, epsilon: 0.010176, episode:  148\n",
      "frames: 260000, reward: -10.900000, loss: 0.000242, epsilon: 0.010171, episode:  148\n",
      "frames: 261000, reward: -10.900000, loss: 0.001355, epsilon: 0.010165, episode:  148\n",
      "frames: 262000, reward: -10.900000, loss: 0.000275, epsilon: 0.010160, episode:  149\n",
      "frames: 263000, reward: -10.900000, loss: 0.000324, epsilon: 0.010154, episode:  149\n",
      "frames: 264000, reward: -10.900000, loss: 0.001632, epsilon: 0.010149, episode:  149\n",
      "frames: 265000, reward: -10.400000, loss: 0.000211, epsilon: 0.010144, episode:  150\n",
      "frames: 266000, reward: -10.400000, loss: 0.002376, epsilon: 0.010140, episode:  150\n",
      "frames: 267000, reward: -10.200000, loss: 0.002369, epsilon: 0.010135, episode:  151\n",
      "frames: 268000, reward: -10.200000, loss: 0.001241, epsilon: 0.010131, episode:  151\n",
      "frames: 269000, reward: -10.000000, loss: 0.002741, epsilon: 0.010126, episode:  152\n",
      "frames: 270000, reward: -10.000000, loss: 0.002905, epsilon: 0.010122, episode:  152\n"
     ]
    }
   ],
   "source": [
    "# Training DRQN in PongNoFrameskip-v4 \n",
    "env = make_atari('PongNoFrameskip-v4')\n",
    "env = wrap_deepmind(env, scale = False, frame_stack=True)\n",
    "\n",
    "gamma = 0.99 # discount factor\n",
    "epsilon_max = 1 # epsilon greedy parameter max\n",
    "epsilon_min = 0.01 # epsilon greedy parameter min\n",
    "eps_decay = 30000 # epsilon greedy parameter decay\n",
    "frames = 1000000 # total training frames\n",
    "USE_CUDA = True # training with gpu\n",
    "learning_rate = 2e-4 # learning rate\n",
    "max_buff = 100000 # maximum buffer size\n",
    "update_tar_interval = 1000 # frames for updating target network\n",
    "batch_size = 32 \n",
    "max_seq = 8\n",
    "\n",
    "print_interval = 1000 \n",
    "log_interval = 1000\n",
    "learning_start = 5000 # 10000\n",
    "\n",
    "action_space = env.action_space\n",
    "action_dim = env.action_space.n\n",
    "state_dim = env.observation_space.shape[0]\n",
    "state_channel = env.observation_space.shape[2]\n",
    "agent = DRQNAgent(in_channels = state_channel, action_space= action_space, USE_CUDA = USE_CUDA, lr = learning_rate, max_seq = max_seq, batch_size = batch_size)\n",
    "frame = env.reset()\n",
    "\n",
    "episode_reward = 0\n",
    "all_rewards = []\n",
    "losses = []\n",
    "episode_num = 0\n",
    "# tensorboard\n",
    "summary_writer = SummaryWriter(log_dir = \"DRQN\", comment= \"good_makeatari\")\n",
    "\n",
    "# e-greedy decay\n",
    "epsilon_by_frame = lambda frame_idx: epsilon_min + (epsilon_max - epsilon_min) * math.exp(\n",
    "            -1. * frame_idx / eps_decay)\n",
    "# plt.plot([epsilon_by_frame(i) for i in range(10000)])\n",
    "\n",
    "hidden = None # initialize hidden state to None\n",
    "for i in range(frames):\n",
    "    epsilon = epsilon_by_frame(i)\n",
    "    state_tensor = agent.observe(frame)\n",
    "    action, hidden = agent.act(state_tensor, hidden, epsilon)\n",
    "    \n",
    "    next_frame, reward, done, _ = env.step(action)\n",
    "    \n",
    "    episode_reward += reward\n",
    "    agent.rec_memory_buffer.push(frame, action, reward, next_frame, done)\n",
    "    frame = next_frame\n",
    "    \n",
    "    loss = 0\n",
    "    if agent.rec_memory_buffer.size() >= learning_start:\n",
    "        loss = agent.learn_from_experience(batch_size)\n",
    "        losses.append(loss)\n",
    "\n",
    "    if i % print_interval == 0:\n",
    "        print(\"frames: %5d, reward: %5f, loss: %4f, epsilon: %5f, episode: %4d\" % (i, np.mean(all_rewards[-10:]), loss, epsilon, episode_num))\n",
    "        summary_writer.add_scalar(\"Temporal Difference Loss\", loss, i)\n",
    "        summary_writer.add_scalar(\"Mean Reward\", np.mean(all_rewards[-10:]), i)\n",
    "        summary_writer.add_scalar(\"Epsilon\", epsilon, i)\n",
    "        \n",
    "    if i % update_tar_interval == 0:\n",
    "        agent.DRQN_target.load_state_dict(agent.DRQN.state_dict())\n",
    "    \n",
    "    if done:\n",
    "\n",
    "        frame = env.reset()\n",
    "        hidden = None # reset hidden to None\n",
    "        all_rewards.append(episode_reward)\n",
    "        episode_reward = 0\n",
    "        episode_num += 1\n",
    "        avg_reward = float(np.mean(all_rewards[-100:]))\n",
    "\n",
    "summary_writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'clear_output' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-2-a68963a494d0>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;31m##### def plot_training(frame_idx, rewards, losses):\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 2\u001b[1;33m \u001b[0mclear_output\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;32mTrue\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      3\u001b[0m \u001b[0mplt\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mfigure\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mfigsize\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m20\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;36m5\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[0mplt\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0msubplot\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m131\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[0mplt\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtitle\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'frame %s. reward: %s'\u001b[0m \u001b[1;33m%\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mframe_idx\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmean\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mrewards\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;33m-\u001b[0m\u001b[1;36m10\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mNameError\u001b[0m: name 'clear_output' is not defined"
     ]
    }
   ],
   "source": [
    "##### def plot_training(frame_idx, rewards, losses):\n",
    "clear_output(True)\n",
    "plt.figure(figsize=(20,5))\n",
    "plt.subplot(131)\n",
    "plt.title('frame %s. reward: %s' % (frame_idx, np.mean(rewards[-10:])))\n",
    "plt.plot(rewards)\n",
    "plt.subplot(132)\n",
    "plt.title('loss')\n",
    "plt.plot(losses)\n",
    "plt.show()\n",
    "\n",
    "plot_training(i, all_rewards, losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
